# -*- coding: UTF-8 -*-
import math
import copy
from collections import namedtuple
import numpy as np
from numpy.lib.arraysetops import unique
from transform import *

# 设置Numpy打印精度
np.set_printoptions(precision=3, suppress=True)

class Ti5Kinematic:
	'''钛虎(Ti5Robot)协作机械臂运动学'''
	# 关节角度下限, 单位: 弧度
	JOINT_ANGLE_LOWERB = np.radians(np.float32([-180.0,  -265.0, -150.0, -180.0, -175.0, -175.0]))
	# 关节角度上限制, 单位: 弧度
	JOINT_ANGLE_UPPERB = np.radians(np.float32([180.0, 85.0, 150.0, 180.0, 175.0, 175.0]))
	# 默认关节角度
	JOINT_ANGLE_DEFAULT = [0 for i in range(6)]
	# 关节编号
	JOINT1 = 0
	JOINT2 = 1
	JOINT3 = 2
	JOINT5 = 3
	JOINT6 = 4
	JOINT8 = 5
	# DH表 连杆长度定义, 单位mm
	D1 = 175
	A2 = 300
	A3 = 300
	A6 = 96
	# 预定义常量
	PI = np.pi
	MINUS_PI = - np.pi
	PI_M_2 = np.pi * 2.0 # pi*2
	PI_D_2 = np.pi / 2.0 # pi/2
	MIMUS_PI_D_2 = -1.0 * PI_D_2 # -pi/2
	
	def __init__(self):
		'''初始化'''
		# 当前关节角度赋值为默认角度
		self.cur_joint_angle = np.copy(self.JOINT_ANGLE_DEFAULT)
		
	def forward_kinematic_v1(self, joint_angle):
		'''正向运动学-腕关节'''
		# 提取关节角度
		theta1, theta2, theta3, theta5, theta6, theta8 = joint_angle
		# 关节与关节之间的变换
		T01 = Transform.dhmat(0, 0, theta1, self.D1)
		T12 = Transform.dhmat(self.MIMUS_PI_D_2, 0, theta2, 0)
		T23 = Transform.dhmat(0, self.A2, theta3, 0)
		T34 = Transform.dhmat(0, self.A3, self.MIMUS_PI_D_2, 0)
		T45 = Transform.dhmat(self.MIMUS_PI_D_2, 0, theta5, 0)
		T56 = Transform.dhmat(self.PI_D_2, 0, theta6+self.PI_D_2, 0)
		T67 = Transform.dhmat(0, self.A6, self.MIMUS_PI_D_2, 0)
		T78 = Transform.dhmat(self.MIMUS_PI_D_2, 0, theta8, 0)
		# 计算机械臂基坐标系到腕关节坐标系的变换
		T02 = T01.dot(T12)
		T03 = T02.dot(T23)
		T04 = T03.dot(T34)
		T05 = T04.dot(T45)
		T06 = T05.dot(T56)
		T07 = T06.dot(T67)
		T08 = T07.dot(T78)
		# 创建位姿对象
		pose = Pose()
		pose.set_transform_matrix(T08)
		return pose, [T01, T02, T03, T04, T05, T06, T07, T08]

	def forward_kinematic_v2(self, joint_angle):
		'''正向运动学-腕关节'''
		# 提取关节角度
		theta1, theta2, theta3, theta5, theta6, theta8 = joint_angle
		# 预先计算三角函数
		s1, s2, s3, s5, s6, s8 = np.sin(joint_angle)
		c1, c2, c3, c5, c6, c8 = np.cos(joint_angle)
		theta23 = theta2 + theta3
		s23, c23 = np.sin(theta23), np.cos(theta23)
		# 求解旋转矩阵
		r11 = - c1*s23*(s5*s8 - c5*c6*c8) \
	  		+ s1*(s5*c6*c8 + c5*s8) \
			+ c1*c23*s6*c8
		r12 = c1*s23*(s5*c8 + c5*c6*s8) \
			-s1*(s5*c6*s8 - c5*c8) \
			- c1*c23*s6*s8
		r13 = - s1*s5*s6 - c1*s23*c5*s6 \
	  		+ c1*c23*c6 
		r21 = - s1*s23*(s5*s8 - c5*c6*c8) \
			- c1*(s5*c6*c8 + c5*s8) \
			+ s1*c23*s6*c8
		r22 = - s1*s23*(s5*c8 + c5*c6*s8) \
			+ c1*(s5*c6*s8 - c5*c8) \
			- s1*c23*s6*s8
		r23 = - s1*s23*c5*s6 + s1*c23*c6 \
	  		+ c1*s5*s6
		r31 = - c23*(s5*s8 - c5*c6*c8) \
			- s23*s6*c8
		r32 = - c23*(s5*c8 + c5*c6*s8) + s23*s6*s8
		r33 = - c23*c5*s6 - s23*c6
		# 临时变量
		A3_ADD_A6_M_C6 = (self.A3 + self.A6*c6)
		# 求解坐标
		x8 = self.A2*c1*c2 - self.A6*s1*s5*s6 \
			- self.A6*c1*s23*c5*s6 \
	   		+ A3_ADD_A6_M_C6*c1*c23
		y8 = self.A2*s1*c2 - self.A6*s1*s23*c5*s6 \
			+ self.A6*c1*s5*s6 \
	   		+ A3_ADD_A6_M_C6*s1*c23
		z8 = - self.A2*s2 - self.A6*c23*c5*s6 \
			+ self.D1 - A3_ADD_A6_M_C6*s23

		# 创建位姿对象
		pose = Pose()
		# 构造旋转矩阵
		rmat = np.float32([
			[r11, r12, r13],
			[r21, r22, r23],
			[r31, r32, r33]])
		pose.set_rotation_matrix(rmat)
		pose.set_position(x8, y8, z8)
		return pose

	def angle_unique(self, angle_list):
		'''得到不重复的角度列表'''
		return np.unique(np.round(np.float32(angle_list), 8))
	
	def is_joint_angle_legal(self, theta, joint_index):
		'''判断关节角度是否合法'''
		lowerb = self.JOINT_ANGLE_LOWERB[joint_index]
		upperb = self.JOINT_ANGLE_UPPERB[joint_index]
		is_legal = theta >= lowerb and theta <= upperb
		angle = np.degrees(theta)
		# print(f"joint{joint_index+1} value: {angle:.1f} legal: {is_legal}")
		return is_legal

	def joint_angle_filter(self, theta_list, joint_index):
		'''关节角度过滤器'''
		# 过滤唯一的候选角度
		theta_list = np.unique(np.float64(theta_list))
		# 筛选后的候选关节角度
		theta_list_filter = []
		for theta in theta_list:
			if self.is_joint_angle_legal(theta, joint_index):
				theta_list_filter.append(theta)
		return theta_list_filter

	def inverse_kinematic(self, pose, last_joint_angle = None, is_debug=False):
		'''逆向运动学'''
		# 上一次的关节角度
		if last_joint_angle is None:
			last_joint_angle = self.JOINT_ANGLE_DEFAULT		
		# 关节边界值
		lowerb  = self.JOINT_ANGLE_LOWERB
		upperb = self.JOINT_ANGLE_UPPERB
		# 候选关节角度
		candi_joint_angle_list = []
		# 提取腕关节的坐标
		x8, y8, z8 = pose.get_position()
		# 旋转矩阵
		R08 = pose.get_rotation_matrix()
		r11, r12, r13, r21, r22, r23, r31, r32, r33 = R08.reshape(-1)
		
  		#  求解P06:  关节6坐标系原点在基坐标系下的坐标
		p08 = np.float32([x8, y8, z8])
		z8_unit = np.float32([r13, r23, r33]) # 关节8坐标系Z轴单位向量	
		p06 = p08 - z8_unit * self.A6         # 关节6坐标系原点在基坐标系下的坐标
		x6, y6, z6 = p06
		# print(f"x6: {x6:.3f}, y6: {y6:.3f}, z6: {z6:.3f}")
		# 判断x6与y6是否同时为0
		# 求解关节1的角度
		theta1_list = None
		if x6 == 0.0 and y6 == 0.0:
			theta1_1 = last_joint_angle[self.JOINT1]
			theta1_list = [theta1_1]
		else:
			theta1_1 = np.arctan2(y6, x6)
			# print(f"theta1: {math.degrees(theta1_1):.5f}")
			theta1_2 = theta1_1 + self.PI
			if theta1_2 > np.pi:
				theta1_2 -= self.PI_M_2
			theta1_list = [theta1_1, theta1_2]
		# 求解关节3的角度
		for theta1 in self.joint_angle_filter(theta1_list,\
	  			self.JOINT1):
			c1 = np.cos(theta1)
			s1 = np.sin(theta1)
			# 求解关节6在关节1坐标系下的坐标
			x_16 = c1*x6 + s1*y6
			z_16 = z6 - self.D1
			# 定义b1, b2 
			b1 = x_16
			b2 = -z_16
			c3 = (b1**2 + b2**2 - self.A2**2 - self.A3**2) / (2*self.A2*self.A3)
			# 数值原因 cos(theta3)可能会再1.0附近
			# 为了被判定为非法, 因此做下特殊处理
			c3 = np.round(c3, 10)
			if np.abs(c3) > 1.0:
				# 非法cos值, 工作区达不到
				# print(f"cos(theta3) = {c3}")
				continue
			s3_abs = np.sqrt(1 - c3**2)
			theta3_1 = np.arctan2(s3_abs, c3)
			theta3_2 = np.arctan2(-s3_abs, c3)
			# 求解theta2
			for theta3 in self.joint_angle_filter(\
				[theta3_1, theta3_2], self.JOINT3):
				s3 = np.sin(theta3)
				# print(f"theta3 = {theta3}")
				K1 = self.A2 + self.A3 * c3
				K2 = self.A3 * s3
				theta2 = np.arctan2(b2, b1) - np.arctan2(K2, K1)
				# 关节角度2的范围有点特殊 [-265.0, 85.0]
				if theta2 < self.JOINT_ANGLE_LOWERB[self.JOINT2]:
					theta2 += self.PI_M_2
				elif theta2 > self.JOINT_ANGLE_UPPERB[self.JOINT2]:
					theta2 -= self.PI_M_2
				# 验证theta2的合法性
				if not self.is_joint_angle_legal(theta2, self.JOINT2):
					continue
				# 求解关节8在关节4下的位姿^{4}_{8}R
				s2, c2 = np.sin(theta2), np.cos(theta2)
				theta23 = theta2 + theta3
				s23, c23 = np.sin(theta23), np.cos(theta23)
				R40 = np.float32([
					[c1*s23, s1*s23, c23],
					[c1*c23, s1*c23, -s23],
			  		[-s1, c1, 0]])
				R48 = np.dot(R40, R08)
				# 提取R48里面的元素
				q11, q12, q13, q21, q22, q23, q31, q32, q33 = R48.reshape(-1)
				# 求解关节6的角度theta6
				sin_theta6_abs = np.sqrt(1 - q23*q23)
				theta6_1 = np.arctan2(sin_theta6_abs, q23) 
				theta6_2 = np.arctan2(-sin_theta6_abs, q23)
				for theta6 in self.joint_angle_filter(\
					[theta6_1, theta6_2], self.JOINT6):
					theta5 = None
					theta8 = 0
					theta58 = None
					s6 = np.sin(theta6)
					if s6 == 0:
						# 出现万向锁的问题
						if theta6 == 0:
							theta5 = last_joint_angle[self.JOINT5]
							theta58 = np.arctan2(-q12, q11)
						else:
							theta5 = last_joint_angle[self.JOINT5]
							theta58 = np.arctan2(q31, -q32)
						theta8 = theta58 - theta5
					else:
						# 普通情况
						if s6 > 0:
							# theta6 > 0
							theta5 = np.arctan2(q33, -q13)
							theta8 = np.arctan2(-q22, q21)
						else:
							# theta6 < 0
							theta5 = np.arctan2(-q33, q13)
							theta8 = np.arctan2(q22, -q21)
					# 验证关节5角度的合法性
					if not self.is_joint_angle_legal(theta5, self.JOINT5):
						continue
					# 验证关节8角度的合法性
					if not self.is_joint_angle_legal(theta8, self.JOINT8):
						continue
					candi_joint_angle_list.append([theta1, theta2, theta3, theta5, theta6, theta8])
		return candi_joint_angle_list

if __name__=="__main__":
    s = Ti5Kinematic()

    while True :
        lst=list()

        pose = Pose()
        print("Please input XYZ and RPY:       ")
        for i in range(0,6):
            ele=float(input())
            lst.append(ele)

        print(str(lst))
        pose.set_position(lst[0]*70,lst[1]*70,lst[2]*70)
        pose.set_euler_angle(lst[3],lst[4],lst[5])
        res = s.inverse_kinematic(pose)
        for i in res:
            
            print(str(i))


